from sklearn.base import BaseEstimator, TransformerMixin
from pathlib import Path
from datetime import datetime

import pandas as pd
import numpy as np
import json
import re

class MultiOptionProcessor(BaseEstimator, TransformerMixin):
    """
    多功能多选项处理器
    功能：
    1. 多选项字符串拆分与编码
    2. 高频选项选择与OTHER类别处理
    3. 训练/测试特征一致性
    4. 元数据保存与加载
    5. 安全特征命名
    6. 与scikit-learn Pipeline集成

    参数:
    max_features: 每列最多保留的高频选项数 (default=10)
    min_frequency: 选项最低出现频率 (default=0.01)
    other_category: 是否添加OTHER类别 (default=True)
    save_mappings: 是否保存特征映射 (default=False)
    split_str: 字符串切分字符(default=",")
    mapping_dir: 特征映射保存目录 (default="feature_mappings")
    """

    def __init__(self, max_features=10, min_frequency=0.01, other_category=True, save_mappings=False,split_str=",",
                 mapping_dir="feature_mappings"):
        self.max_features = max_features
        self.min_frequency = min_frequency
        self.other_category = other_category
        self.save_mappings = save_mappings
        self.split_str = split_str
        self.mapping_dir = Path(mapping_dir)
        self.feature_mappings = {}
        self.feature_names_ = []
        self.is_fitted = False

    def fit(self, X, columns=None):
        """
        在数据上拟合处理器
        X: pandas DataFrame
        y: 忽略 (用于scikit-learn兼容)
        columns: 指定处理的列 (默认处理所有字符串列)
        """

        if not isinstance(X, pd.DataFrame):
            raise ValueError("输入必须是pandas DataFrame")

        # 确定要处理的列
        self.columns = columns or self._detect_multivalue_columns(X)
        self.feature_mappings = {}
        self.feature_names_ = []

        # 创建映射目录
        if self.save_mappings:
            self.mapping_dir.mkdir(exist_ok=True, parents=True)

        for col in self.columns:
            mapping = self._create_feature_mapping(X, col)
            self.feature_mappings[col] = mapping
            self.feature_names_.extend(mapping['feature_names'])

            # 保存映射
            if self.save_mappings:
                self._save_mapping(col, mapping)

        self.is_fitted = True
        return self

    def transform(self, X):
        """转换数据"""
        if not self.is_fitted:
            raise RuntimeError("请先调用fit()方法")

        if not isinstance(X, pd.DataFrame):
            X = pd.DataFrame(X)

        # 创建结果副本
        result = X.copy()

        for col, mapping in self.feature_mappings.items():
            if col not in result.columns:
                continue

            # 处理选项
            processed = self._process_column(result[col], mapping)

            # 添加新特征
            for feature in mapping['feature_names']:
                result[feature] = processed[feature]

            # 移除原始列
            result = result.drop(col, axis=1)

        return result

    def process_test_set(self, X_test, handle_unknown='other'):
        """专门处理测试集的便捷方法（自动加载映射）"""
        if not self.is_fitted and self.save_mappings:
            self.load_mappings()
        return self.transform(X_test)

    def load_mappings(self, mapping_dir=None):
        """加载保存的特征映射"""
        dir_path = Path(mapping_dir) if mapping_dir else self.mapping_dir
        self.feature_mappings = {}
        self.feature_names_ = []

        for file in dir_path.glob("*_mapping.json"):
            col = file.name.replace("_mapping.json", "")
            with open(file, 'r') as f:
                mapping = json.load(f)
                self.feature_mappings[col] = mapping
                self.feature_names_.extend(mapping['feature_names'])

        self.is_fitted = True
        return self

    def get_feature_names(self):
        """获取生成的特征名"""
        return self.feature_names_

    def get_feature_names_out(self, input_features=None):
        """兼容scikit-learn的get_feature_names_out"""
        return np.array(self.feature_names_)

    def _detect_multivalue_columns(self, df):
        """检测可能是多选项的列"""
        candidate_cols = []
        for col in df.columns:
            if pd.api.types.is_string_dtype(df[col]):
                # 检查是否有逗号分隔的值
                sample = df[col].dropna().sample(min(100, len(df)))
                if any(',' in str(x) for x in sample):
                    candidate_cols.append(col)
        return candidate_cols

    def _create_feature_mapping(self, df, col):
        """为单列创建特征映射"""
        # 预处理数据
        cleaned = self._clean_and_split(df[col])

        # 统计选项频率
        all_items = [item for sublist in cleaned for item in sublist]
        if not all_items:
            return {'feature_names': [], 'options': []}

        item_counts = pd.Series(all_items).value_counts(normalize=True)
        total_rows = len(df)

        # 确定高频选项
        valid_items = item_counts[item_counts >= self.min_frequency].index
        top_items = valid_items[:min(self.max_features, len(valid_items))]

        # 创建特征名
        feature_names = [self._safe_feature_name(col, item) for item in top_items]

        # 添加OTHER类别
        if self.other_category and (len(valid_items) > len(top_items)):
            other_name = self._safe_feature_name(col, "OTHER")
            feature_names.append(other_name)

        # 构建映射
        mapping = {
            'top_options': list(top_items),
            'all_options': list(item_counts.index),
            'has_other': self.other_category and (len(valid_items) > len(top_items)),
            'feature_names': feature_names,
            'created_at': datetime.now().isoformat(),
            'total_rows': total_rows,
            'config': {
                'max_features': self.max_features,
                'min_frequency': self.min_frequency,
                'other_category': self.other_category
            }
        }

        return mapping

    def _process_column(self, series, mapping):
        """处理单列数据"""
        # 预处理数据
        cleaned = self._clean_and_split(series)

        # 创建结果容器
        result = pd.DataFrame(index=series.index)

        # 添加高频选项特征
        for item in mapping['top_options']:
            feature = self._safe_feature_name(series.name, item)
            result[feature] = [1 if item in items else 0 for items in cleaned]

        # 添加OTHER特征
        if mapping.get('has_other', False):
            other_feature = self._safe_feature_name(series.name, "OTHER")
            top_set = set(mapping['top_options'])
            result[other_feature] = [
                1 if any(item not in top_set for item in items) else 0
                for items in cleaned
            ]

        return result

    def _clean_and_split(self, series):
        """清理并分割选项"""
        return (
            series
            .fillna('')
            .astype(str)
            .str.upper()
            .str.replace(r'\s+', '', regex=True)  # 移除所有空白
            .apply(lambda x: [item for item in x.split(self.split_str) if item])
        )

    def _safe_feature_name(self, col, item):
        """生成安全的特征名"""
        # 移除特殊字符，保留字母、数字和下划线
        safe_col = re.sub(r'[^\w]', '_', col)
        safe_item = re.sub(r'[^\w]', '_', item)
        return f"{safe_col}_{safe_item}"

    def _save_mapping(self, col, mapping):
        """保存特征映射到文件"""
        file_path = self.mapping_dir / f"{col}_mapping.json"
        with open(file_path, 'w') as f:
            json.dump(mapping, f, indent=2)

        # 保存元数据
        metadata = {
            'processor_version': '1.0.0',
            'created_at': datetime.now().isoformat(),
            'processed_columns': list(self.feature_mappings.keys()),
            'total_features': len(self.feature_names_),
            'config': {
                'max_features': self.max_features,
                'min_frequency': self.min_frequency,
                'other_category': self.other_category
            }
        }
        with open(self.mapping_dir / "metadata.json", 'w') as f:
            json.dump(metadata, f, indent=2)

    def __repr__(self):
        return (f"MultiOptionProcessor(max_features={self.max_features}, "
                f"min_frequency={self.min_frequency}, "
                f"other_category={self.other_category})")